import gradio as gr
import numpy as np
import cv2
from sklearn.cluster import KMeans
from PIL import Image
import warnings
warnings.filterwarnings('ignore')

def extract_palette(image, n_colors=5):
    """
    使用 KMeans 算法从图像中提取主色调调色板
    """
    try:
        # 转换图像为 numpy 数组
        if isinstance(image, str):
            # 如果是文件路径
            img = cv2.imread(image)
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        else:
            # 如果是上传的图像对象
            img = np.array(image)
            # 如果图像有透明度通道，只取 RGB
            if img.shape[2] == 4:
                img = img[:, :, :3]
        
        # 重塑图像为像素点列表
        pixels = img.reshape(-1, 3)
        
        # 使用 KMeans 聚类找到主要颜色
        kmeans = KMeans(n_clusters=n_colors, random_state=42, n_init=10)
        kmeans.fit(pixels)
        
        # 获取聚类中心
        colors = kmeans.cluster_centers_.astype(int)
        
        return colors
        
    except Exception as e:
        print(f"Error in extract_palette: {e}")
        return None

def rgb_to_hex(rgb):
    """将 RGB 颜色转换为十六进制格式"""
    return '#{:02x}{:02x}{:02x}'.format(rgb[0], rgb[1], rgb[2])

def create_palette_image(colors, palette_size=(400, 80)):
    """
    创建调色板图像
    """
    try:
        palette_width, palette_height = palette_size
        color_width = palette_width // len(colors)
        
        # 创建调色板图像
        palette_img = np.zeros((palette_height, palette_width, 3), dtype=np.uint8)
        
        for i, color in enumerate(colors):
            start_x = i * color_width
            end_x = (i + 1) * color_width
            if i == len(colors) - 1:  # 最后一个颜色填满剩余空间
                end_x = palette_width
            palette_img[:, start_x:end_x] = color
        
        return palette_img
    except Exception as e:
        print(f"Error in create_palette_image: {e}")
        return None

def process_image(image):
    """
    处理上传的图像并返回调色板结果
    """
    if image is None:
        return None, "请上传一张图片"
    
    try:
        # 提取调色板
        colors = extract_palette(image)
        
        if colors is None:
            return None, "无法处理该图像"
        
        # 创建调色板图像
        palette_img = create_palette_image(colors)
        
        if palette_img is None:
            return None, "无法生成调色板"
        
        # 转换为 PIL 图像用于显示
        palette_pil = Image.fromarray(palette_img)
        
        # 生成十六进制颜色码和显示文本
        hex_codes = [rgb_to_hex(color) for color in colors]
        hex_text = "\n".join([f"{hex_code}" for hex_code in hex_codes])
        
        return palette_pil, hex_text
        
    except Exception as e:
        print(f"Error processing image: {e}")
        return None, f"处理图像时出错: {str(e)}"

def create_interface():
    """创建 Gradio 界面"""
    
    with gr.Blocks(
        title="KMeans调色板提取器",
        theme=gr.themes.Soft(),
        css="""
        .gradio-container {
            max-width: 900px !important;
        }
        .color-hex {
            font-family: 'Courier New', monospace;
            font-weight: bold;
            padding: 8px;
            margin: 4px;
            border-radius: 4px;
            text-align: center;
        }
        """
    ) as demo:
        
        gr.Markdown(
            """
            # 🎨 KMeans调色板提取器
            **上传图像获取其主色调调色板（5个颜色）以及每一个颜色的十六进制表示形式**
            """
        )
        
        with gr.Row():
            with gr.Column(scale=1):
                # 图像上传区域
                image_input = gr.Image(
                    label="上传图像 (Image)",
                    type="pil",
                    height=300,
                    sources=["upload", "clipboard"]
                )
                
                with gr.Row():
                    clear_btn = gr.Button("🔄 Clear", variant="secondary", size="sm")
                    submit_btn = gr.Button("🚀 Submit", variant="primary", size="sm")
            
            with gr.Column(scale=1):
                # 调色板显示区域
                palette_output = gr.Image(
                    label="提取的调色板",
                    height=120,
                    show_download_button=True,
                    show_share_button=True
                )
                
                # 十六进制颜色码显示
                hex_output = gr.Textbox(
                    label="色彩十六进制码",
                    lines=6,
                    max_lines=6,
                    placeholder="调色板颜色将显示在这里...",
                    show_copy_button=True
                )
        
        # 状态显示
        status = gr.Textbox(
            label="状态",
            value="就绪",
            interactive=False,
            max_lines=1
        )
        
        # 示例图像
        gr.Markdown("### 🖼️ 示例图像")
        gr.Examples(
            examples=[
                ["examples/sample1.jpg"],  # 你可以替换为自己的示例图像路径
                ["examples/sample2.jpg"],
                ["examples/sample3.jpg"]
            ],
            inputs=image_input,
            outputs=[palette_output, hex_output],
            fn=process_image,
            cache_examples=False,
            label="点击示例图像快速测试"
        )
        
        def update_status(message):
            return message
        
        # 按钮事件绑定
        submit_btn.click(
            fn=update_status,
            inputs=gr.Textbox(value="处理中...", visible=False),
            outputs=status
        ).then(
            fn=process_image,
            inputs=image_input,
            outputs=[palette_output, hex_output]
        ).then(
            fn=update_status,
            inputs=gr.Textbox(value="处理完成！", visible=False),
            outputs=status
        )
        
        clear_btn.click(
            fn=lambda: [None, "", "已清除"],
            inputs=[],
            outputs=[image_input, hex_output, status]
        )
        
        # 自动处理上传的图像
        image_input.upload(
            fn=update_status,
            inputs=gr.Textbox(value="检测到图像上传，处理中...", visible=False),
            outputs=status
        ).then(
            fn=process_image,
            inputs=image_input,
            outputs=[palette_output, hex_output]
        ).then(
            fn=update_status,
            inputs=gr.Textbox(value="处理完成！", visible=False),
            outputs=status
        )
    
    return demo

# 创建并启动应用
if __name__ == "__main__":
    demo = create_interface()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True,
        inbrowser=True  # 自动在浏览器中打开
    )